package app.framework.security.csrf;

import app.framework.constant.AuthConstants;
import cn.hutool.core.codec.Base64;
import cn.hutool.core.util.ArrayUtil;
import cn.hutool.crypto.digest.Digester;
import jasmine.framework.common.exception.InvalidPropertyException;
import jasmine.framework.common.util.CheckUtil;
import jasmine.framework.common.util.StringUtil;
import jasmine.framework.context.WithContext;

import javax.servlet.http.Cookie;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpSession;
import java.util.UUID;

/**
 * @author mh.z
 */
public class CsrfTokenUtil implements WithContext {
    private static Digester digester;

    public static void initUtil(Digester digester) {
        CsrfTokenUtil.digester = digester;
    }

    /**
     * 检查 CSRF 令牌
     *
     * @param request
     * @return
     */
    public static boolean checkCsrfToken(HttpServletRequest request) {
        CheckUtil.notNull(request, "request null");

        // 获取请求中记录的 CSRF 令牌
        String requestCsrfToken = request.getHeader(AuthConstants.CSRF_TOKEN_REQUEST_HEAD_NAME);
        if (StringUtil.isEmpty(requestCsrfToken)) {
            return false;
        }

        boolean result = false;
        HttpSession session = request.getSession(false);

        if (session != null) {
            // 获取会话中记录的 CSRF 令牌
            String sessionCsrfToken = (String) session.getAttribute(AuthConstants.CSRF_TOKEN_SESSION_ATTR_NAME);

            result = StringUtil.equals(requestCsrfToken, sessionCsrfToken);
        }

        if (!result) {
            // 获取"记住我"令牌
            String rememberMeToken = obtainRememberMeToken(request);

            if (StringUtil.isNotEmpty(rememberMeToken)) {
                // 加密”记住我"令牌
                String encodedRememberMeToken = encodeRememberMeToken(rememberMeToken);

                result = StringUtil.equals(requestCsrfToken, encodedRememberMeToken);
            }
        }

        return result;
    }

    /**
     * 设置 CSRF 令牌
     *
     * @param request
     * @param response
     */
    public static void setCsrfToken(HttpServletRequest request, HttpServletResponse response) {
        CheckUtil.notNull(request, "request null");
        CheckUtil.notNull(response, "response null");

        String rememberMeToken = obtainRememberMeToken(request);
        String csrfToken;

        if (StringUtil.isNotEmpty(rememberMeToken)) {
            csrfToken = encodeRememberMeToken(rememberMeToken);
        } else {
            csrfToken = UUID.randomUUID().toString();
            csrfToken = Base64.encodeUrlSafe(csrfToken);
        }

        HttpSession session = request.getSession(true);
        session.setAttribute(AuthConstants.CSRF_TOKEN_SESSION_ATTR_NAME, csrfToken);

        Cookie csrfCookie = new Cookie(AuthConstants.CSRF_TOKEN_COOKIE_NAME, csrfToken);
        String contextPath = request.getContextPath();
        String cookiePath = (contextPath.length() > 0)
                ? contextPath : "/";
        csrfCookie.setPath(cookiePath);
        // 添加 CSRF 令牌到 cookie 中
        response.addCookie(csrfCookie);
    }

    /**
     * 加密”记住我"令牌
     *
     * @param rememberMeToken
     * @return
     */
    public static String encodeRememberMeToken(String rememberMeToken) {
        CheckUtil.notNull(rememberMeToken, "rememberMeToken null");

        if (digester == null) {
            throw new InvalidPropertyException("CsrfTokenUtil.digester null", null);
        }

        byte[] bytes = digester.digest(rememberMeToken);
        String result = Base64.encodeUrlSafe(bytes);

        return result;
    }

    /**
     * 查找”记住我"令牌
     *
     * @param request
     * @return
     */
    public static String obtainRememberMeToken(HttpServletRequest request) {
        CheckUtil.notNull(request, "request null");

        String rememberMeToken = null;
        HttpSession session = request.getSession(false);

        if (session != null) {
            rememberMeToken = (String) session.getAttribute(AuthConstants.REMEMBER_ME_SESSION_ATTR_NAME);
        }

        if (StringUtil.isEmpty(rememberMeToken)) {
            Cookie[] cookies = request.getCookies();

            if (ArrayUtil.isNotEmpty(cookies)) {
                for (Cookie cookie : cookies) {

                    if (StringUtil.equals(AuthConstants.REMEMBER_ME_COOKIE_NAME, cookie.getName())) {
                        rememberMeToken = cookie.getValue();
                        break;
                    }
                }
            }
        }

        return rememberMeToken;
    }

}
